package covertjava.bytecode;

import java.io.IOException;
import java.util.regex.Pattern;

import org.apache.bcel.*;
import org.apache.bcel.classfile.*;
import org.apache.bcel.generic.*;

/**
 * <p>Instruments bytecode methods by adding an invocation of a custom method</p>
 * <p>Copyright: Copyright (c) 2004 Sams Publishing</p>
 * @author Alex Kalinovsky
 * @version 1.0
 */
public class MethodInstrumentor {

    public static void main(String[] args) throws IOException {
        if (args.length != 2) {
            System.out.println("Syntax: MethodInstrumentor <full class name> <method name pattern>");
            System.exit(1);
        }
        JavaClass cls = Repository.lookupClass(args[0]);
        MethodInstrumentor instrumentor = new MethodInstrumentor();
        instrumentor.instrumentWithInvocationRegistry(cls, args[1]);
        cls.dump("new_" + cls.getClassName() + ".class");
    }

    /**
     * Adds invocation of InvocationRegistry to all methods whose name satisfies the pattern
     */
    public void instrumentWithInvocationRegistry(JavaClass cls, String methodPattern) {
        ConstantPoolGen constants = new ConstantPoolGen(cls.getConstantPool());
        Method[] methods = cls.getMethods();

        for (int i = 0; i < methods.length; i++) {
            // Instrument all methods that match the given criteria
            if (Pattern.matches(methodPattern, methods[i].getName())) {
                methods[i] = instrumentMethod(cls, constants, methods[i]);
            }
        }
        cls.setMethods(methods);
        cls.setConstantPool(constants.getFinalConstantPool());
    }


    public Method instrumentMethod(JavaClass cls, ConstantPoolGen constants, Method oldMethod) {
        System.out.println("Instrumenting method " + oldMethod.getName());
        MethodGen method = new MethodGen(oldMethod, cls.getClassName(), constants);
        InstructionFactory factory = new InstructionFactory(constants);
        InstructionList instructions = new InstructionList();

        // Append two instructions representing a method call
        instructions.append(new PUSH(constants, method.getName()));
        Instruction invoke = factory.createInvoke(
                "bytecode.InvocationRegistry",
                "methodInvoked",
                Type.VOID,
                new Type[] {new ObjectType("java.lang.String")},
                Constants.INVOKESTATIC
                );
        instructions.append(invoke);

        method.getInstructionList().insert(instructions);
        instructions.dispose();
        return method.getMethod();
    }
}
